use crate::integrations::windows::ffi::*;
use crate::integrations::windows::socket_table_iterator::SocketTableIterator;
use crate::types::*;
use crate::types::error::*;

/// Iterate through sockets information.
pub fn iterate_sockets_info(
    af_flags: AddressFamilyFlags,
    proto_flags: ProtocolFlags,
) -> Result<impl Iterator<Item=Result<SocketInfo, Error>>, Error> {
    let mut iterators = Vec::with_capacity(4);
    match af_flags {
        AddressFamilyFlags::IPV4 => {
            match proto_flags {
                ProtocolFlags::TCP => {
                    iterators.push(SocketTableIterator::new::<MIB_TCPTABLE_OWNER_PID>()?);
                }
                ProtocolFlags::UDP => {
                    iterators.push(SocketTableIterator::new::<MIB_UDPTABLE_OWNER_PID>()?);
                }
            }
        }
        AddressFamilyFlags::IPV6 => {
            match proto_flags {
                ProtocolFlags::TCP => {
                    iterators.push(SocketTableIterator::new::<MIB_TCP6TABLE_OWNER_PID>()?);
                }
                ProtocolFlags::UDP => {
                    iterators.push(SocketTableIterator::new::<MIB_UDP6TABLE_OWNER_PID>()?);
                }
            }
        }
    }

    Ok(iterators.into_iter().flatten())
}

/// Iterate through sockets information. Works on older versions of Windows (like XP and 2003).
pub fn iterate_sockets_info_without_pids(
    proto_flags: ProtocolFlags,
) -> Result<impl Iterator<Item=Result<SocketInfo, Error>>, Error> {
    let mut iterators = Vec::with_capacity(4);
    if proto_flags == ProtocolFlags::TCP {
        iterators.push(SocketTableIterator::new::<MIB_TCPTABLE>()?);
    }
    if proto_flags == ProtocolFlags::UDP {
        iterators.push(SocketTableIterator::new::<MIB_UDPTABLE>()?);
    }

    Ok(iterators.into_iter().flatten())
}
